// sm4.go

package sm4

import (
	"bytes"
	"crypto/cipher"
	"encoding/hex"
	"fmt"
	"github.com/tjfoc/gmsm/sm4"
)

const (
	keyHex = "0123456789abcdef0123456789abcdef" // 16字节的十六进制字符串密钥
	ivHex  = "fedcba98765432100123456789abcdef" // 16字节的十六进制字符串IV
)

// EncryptSM4 使用SM4算法对明文进行加密
func EncryptSM4(plaintext string) (string, error) {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return "", err
	}

	iv, err := hex.DecodeString(ivHex)
	if err != nil {
		return "", err
	}

	data := []byte(plaintext)

	// 使用PKCS7填充
	data = pkcs7Pad(data, sm4.BlockSize)

	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}

	ciphertext := make([]byte, len(data))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, data)

	// 返回加密结果的十六进制字符串
	return hex.EncodeToString(ciphertext), nil
}

// DecryptSM4 使用SM4算法对密文进行解密
func DecryptSM4(ciphertext string) (string, error) {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return "", err
	}

	iv, err := hex.DecodeString(ivHex)
	if err != nil {
		return "", err
	}

	data, err := hex.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}

	decrypted := make([]byte, len(data))
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(decrypted, data)

	// PKCS7填充解除
	decrypted, err = pkcs7Unpad(decrypted)
	if err != nil {
		return "", err
	}

	return string(decrypted), nil
}

// PKCS7填充
func pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS7填充解除
func pkcs7Unpad(data []byte) ([]byte, error) {
	length := len(data)
	unpadding := int(data[length-1])
	if unpadding > length || unpadding == 0 {
		return nil, fmt.Errorf("Invalid padding")
	}
	return data[:length-unpadding], nil
}
